import pandas as pd
import numpy as np
from matplotlib import pyplot as plt


def calculate_bollinger_bands(df, window = 20, no_of_std = 2, column_name = ''):
    '''
    Parameters
    ----------
    df : 股价数据集
    window : 计算的周期，默认20天
    no_of_std : 标准偏差的倍数，默认2
    column_name : 股价数据集选择的列名，默认为第一列
    '''
    if column_name == '':
        column_name = df.columns[0]
    df['Observation']=df[column_name]
    df['RollingMean'] = df[column_name].rolling(window).mean()
    std = df[column_name].rolling(window).std(ddof=0)
    df['UpperBound'] = df['RollingMean'] + no_of_std * std
    df['LowerBound'] = df['RollingMean'] - no_of_std * std
    bb = df[['Observation','RollingMean','UpperBound','LowerBound']]
    bb.plot(title='Observation with Bollinger Bands') #绘制布林带
    return df


def plot_bands_with_signal(df):
    # 绘制交易图表
    # 买入开仓
    long_entry = df.loc[df['signal_long'] == 1]['Observation']
    # 卖出平仓
    long_exit = df.loc[df['signal_long'] == 0]['Observation']
    # 卖出开仓
    short_entry = df.loc[df['signal_short'] == -1]['Observation']
    # 卖出平仓
    short_exit = df.loc[df['signal_short'] == 0]['Observation']

    fig, ax = plt.subplots(1, figsize=(15, 10), sharex=True)
    ax.plot(df['Observation'], label='Observation')
    ax.plot(df['RollingMean'], label='RollingMean')
    ax.plot(df['UpperBound'])
    ax.plot(df['LowerBound'])
    ax.fill_between(df.index, df['UpperBound'], df['LowerBound'],
                    alpha=0.3, label='Bollinger Band')

    ax.scatter(long_entry.index, long_entry, color='r',
               s=100, marker='^', label='Long Entry',
               zorder=10)
    ax.scatter(long_exit.index, long_exit, color='r',
               s=100, marker='x', label='Long Exit',
               zorder=10)
    ax.scatter(short_entry.index, short_entry, color='b',
               s=100, marker='^', label='Short Entry',
               zorder=10)
    ax.scatter(short_exit.index, short_exit, color='b',
               s=100, marker='x', label='Short Exit',
               zorder=10)
    ax.set_title('Bollinger Band Strategy Trading Signals')
    ax.legend()
    plt.show()

def calculate_strategy1_position(df):
    '''
    策略1：均值回归策略
    当收盘价由下向上穿过上轨的时候，做空；然后由上向下穿过中轨的时候，平仓。
    当收盘价由上向下穿过下轨的时候，做多；然后由下向上穿过中轨的时候，平仓。
    '''
    xs = (df.Observation - df.RollingMean) / (df.Observation.shift(1) - df.RollingMean.shift(1))
    df['position'] = np.nan
    df['position'] = np.where(df.Observation<df.LowerBound, 1,df['position'])
    df['position'] = np.where(df.Observation>df.UpperBound ,-1,df['position'])
    df['position'] = np.where(xs<0, 0, df['position'])
    df['position'] = df['position'].ffill().fillna(0)
    #买入开仓
    long_entry_condition = (df['position']==1) & (df['position'].shift(1)!=1)
    df.loc[long_entry_condition,'signal_long'] = 1
    #卖出平仓
    long_exit_condition = (df['position']!=1) &(df['position'].shift(1)==1)
    df.loc[long_exit_condition,'signal_long'] = 0
    #卖出开仓
    short_entry_condition = (df['position']==-1) & (df['position'].shift(1)!=-1)
    df.loc[short_entry_condition,'signal_short'] = -1
    #买入平仓
    short_exit_condition = (df['position']!=-1) & (df['position'].shift(1)==-1)
    df.loc[short_exit_condition,'signal_short'] = 0
    return df

def calculate_strategy2_position(df):
    '''
    -------
    策略2：
    当收盘价由下向上穿过上轨的时候，做多；然后由上向下穿过上轨的时候，平仓。
    当收盘价由上向下穿过下轨的时候，做空；然后由下向上穿过下轨的时候，平仓。
    '''
    #position : 持仓头寸，多仓为1，不持仓为0，空仓为-1
    #siganal : 交易信号，做多为1，平仓为0，做空为-1
    df['position'] = np.where(df.Observation > df.UpperBound, 1, 0)
    df['position'] = np.where(df.Observation < df.LowerBound, -1, df['position'])
    #买入开仓
    long_entry_condition = (df['position']==1) & (df['position'].shift(1)!=1)
    df.loc[long_entry_condition,'signal_long'] = 1
    #卖出平仓
    long_exit_condition = (df['position']!=1) &(df['position'].shift(1)==1)
    df.loc[long_exit_condition,'signal_long'] = 0
    #卖出开仓
    short_entry_condition = (df['position']==-1) & (df['position'].shift(1)!=-1)
    df.loc[short_entry_condition,'signal_short'] = -1
    #买入平仓
    short_exit_condition = (df['position']!=-1) & (df['position'].shift(1)==-1)
    df.loc[short_exit_condition,'signal_short'] = 0
    return df


def calculate_returns(df):
    # 计算市场对数收益率、策略对数收益率、策略超常收益率
    df['market_log_returns'] = np.log(df['Observation']/df['Observation'].shift(1))
    df['strat_log_returns'] = df['position'].shift(1)* df['market_log_returns'] #交易信号出现后第二天交易
    df['abnormal_returns'] = df['strat_log_returns'] - df['market_log_returns']
    # 计算市场累积收益率、策略累积收益率、策略累积超常收益率
    df['market_cum_returns'] = np.exp(df['market_log_returns'].cumsum()) - 1
    df['strat_cum_returns'] = np.exp(df['strat_log_returns'].cumsum()) - 1
    df['abnormal_cum_returns'] = df['strat_cum_returns'] - df['market_cum_returns']
    # 绘图
    ret = df[['market_log_returns','strat_log_returns','abnormal_returns']]
    ret.plot(title='Market,Strategy,and Abnormal Returns')
    cum_returns = df[['market_cum_returns','strat_cum_returns','abnormal_cum_returns']]
    cum_returns.plot(title='Cumulative Returns')
    return df

df = pd.read_csv(r'.\SPY.csv')
df = df.set_index('Date')
df = calculate_bollinger_bands(df, 10, 2, 'Adj Close') #计算布林带

# 策略1
# df1 = calculate_strategy1_position(df.copy())
# df1 = calculate_returns(df1)
# plot_bands_with_signal(df1)

# # 策略2
df2 = calculate_strategy2_position(df.copy())
df2 = calculate_returns(df2)
plot_bands_with_signal(df2)





